"""
13_similarity_algorithms.py
===========================
Graph Similarity Algorithms

This module covers:
1. Node Similarity (Jaccard)
2. K-Nearest Neighbors
3. Cosine Similarity
4. Overlap Similarity
5. Custom similarity functions

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver

try:
    from graphdatascience import GraphDataScience
    GDS_AVAILABLE = True
except ImportError:
    GDS_AVAILABLE = False


class SimilarityAnalyzer:
    """
    Computes similarity between nodes using GDS algorithms.
    
    Similarity Metrics:
    
    ┌─────────────────────────────────────────────────────────────────┐
    │                    SIMILARITY METRICS                           │
    ├─────────────────────────────────────────────────────────────────┤
    │                                                                 │
    │  JACCARD SIMILARITY                                            │
    │  ──────────────────                                             │
    │  |A ∩ B| / |A ∪ B|                                              │
    │                                                                 │
    │  • Measures overlap of neighbor sets                           │
    │  • Range: 0 (no overlap) to 1 (identical)                     │
    │  • Best for: Binary relationships                              │
    │                                                                 │
    │  Example: Patients with same diseases                          │
    │  Patient A: {Diabetes, Hypertension}                           │
    │  Patient B: {Diabetes, Asthma}                                 │
    │  Jaccard = 1/3 = 0.33                                         │
    │                                                                 │
    │  COSINE SIMILARITY                                             │
    │  ─────────────────                                              │
    │  (A · B) / (||A|| × ||B||)                                     │
    │                                                                 │
    │  • Measures angle between vectors                              │
    │  • Range: -1 to 1 (usually 0 to 1)                            │
    │  • Best for: Weighted relationships                           │
    │                                                                 │
    │  OVERLAP SIMILARITY                                            │
    │  ──────────────────                                             │
    │  |A ∩ B| / min(|A|, |B|)                                       │
    │                                                                 │
    │  • Ignores size differences                                    │
    │  • Range: 0 to 1                                               │
    │  • Best for: Subset relationships                              │
    │                                                                 │
    │  K-NEAREST NEIGHBORS (KNN)                                     │
    │  ─────────────────────────                                      │
    │  For each node, find K most similar nodes                      │
    │                                                                 │
    │  • Creates similarity relationships                            │
    │  • Parameters: K (neighbors), similarity metric               │
    │  • Best for: Recommendations, clustering                       │
    │                                                                 │
    └─────────────────────────────────────────────────────────────────┘
    """
    
    def __init__(self, uri: str = None, username: str = None, password: str = None):
        if not GDS_AVAILABLE:
            raise ImportError("graphdatascience package required")
        
        uri = uri or os.getenv("NEO4J_URI", "bolt://localhost:7687")
        username = username or os.getenv("NEO4J_USERNAME", "neo4j")
        password = password or os.getenv("NEO4J_PASSWORD", "password123")
        
        self.gds = GraphDataScience(uri, auth=(username, password))
        self.driver = get_driver()
    
    # =========================================================
    # GRAPH PROJECTION FOR SIMILARITY
    # =========================================================
    
    def project_bipartite_graph(
        self,
        graph_name: str,
        source_label: str,
        target_label: str,
        relationship_type: str
    ) -> Dict[str, Any]:
        """
        Project a bipartite graph for similarity computation.
        
        Bipartite graphs have two node types connected by relationships.
        Example: Patient -[:DIAGNOSED_WITH]-> Disease
        
        Similarity is computed between source nodes based on
        their shared connections to target nodes.
        """
        # Drop if exists
        if self.gds.graph.exists(graph_name)["exists"]:
            self.gds.graph.drop(graph_name)
        
        G, result = self.gds.graph.project(
            graph_name,
            [source_label, target_label],
            {relationship_type: {"orientation": "UNDIRECTED"}}
        )
        
        return {
            "graph": graph_name,
            "nodeCount": result["nodeCount"],
            "relationshipCount": result["relationshipCount"]
        }
    
    # =========================================================
    # NODE SIMILARITY (JACCARD)
    # =========================================================
    
    def node_similarity_stream(
        self,
        graph_name: str,
        top_k: int = 10,
        similarity_cutoff: float = 0.0,
        degree_cutoff: int = 1
    ) -> List[Dict[str, Any]]:
        """
        Compute Node Similarity using Jaccard.
        
        Args:
            graph_name: Name of projected graph
            top_k: Number of similar nodes per node
            similarity_cutoff: Minimum similarity score
            degree_cutoff: Minimum relationships required
            
        Returns:
            List of similar node pairs with scores
        """
        G = self.gds.graph.get(graph_name)
        
        result = self.gds.nodeSimilarity.stream(
            G,
            topK=top_k,
            similarityCutoff=similarity_cutoff,
            degreeCutoff=degree_cutoff
        )
        
        return result.to_dict('records')
    
    def node_similarity_write(
        self,
        graph_name: str,
        write_relationship_type: str = "SIMILAR",
        write_property: str = "score",
        top_k: int = 10,
        similarity_cutoff: float = 0.1
    ) -> Dict[str, Any]:
        """
        Compute Node Similarity and write results to Neo4j.
        
        Creates SIMILAR relationships between similar nodes.
        """
        G = self.gds.graph.get(graph_name)
        
        result = self.gds.nodeSimilarity.write(
            G,
            writeRelationshipType=write_relationship_type,
            writeProperty=write_property,
            topK=top_k,
            similarityCutoff=similarity_cutoff
        )
        
        return {
            "nodesCompared": result["nodesCompared"],
            "relationshipsWritten": result["relationshipsWritten"],
            "computeMillis": result["computeMillis"]
        }
    
    # =========================================================
    # K-NEAREST NEIGHBORS
    # =========================================================
    
    def knn_stream(
        self,
        graph_name: str,
        k: int = 10,
        similarity_metric: str = "JACCARD",
        node_properties: List[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Compute K-Nearest Neighbors.
        
        Args:
            graph_name: Name of projected graph
            k: Number of neighbors per node
            similarity_metric: JACCARD, COSINE, EUCLIDEAN, PEARSON
            node_properties: Properties to use for similarity
            
        KNN can use either:
        - Graph topology (neighbor sets)
        - Node properties (feature vectors)
        """
        G = self.gds.graph.get(graph_name)
        
        config = {
            "topK": k,
            "nodeProperties": node_properties or ["*"]
        }
        
        result = self.gds.knn.stream(G, **config)
        
        return result.to_dict('records')
    
    def knn_mutate(
        self,
        graph_name: str,
        mutate_relationship_type: str = "KNN_SIMILAR",
        mutate_property: str = "score",
        k: int = 10
    ) -> Dict[str, Any]:
        """
        Compute KNN and add relationships to projected graph.
        
        Mutate adds to the in-memory graph without writing to Neo4j.
        Useful for chaining algorithms.
        """
        G = self.gds.graph.get(graph_name)
        
        result = self.gds.knn.mutate(
            G,
            mutateRelationshipType=mutate_relationship_type,
            mutateProperty=mutate_property,
            topK=k
        )
        
        return {
            "nodesCompared": result["nodesCompared"],
            "relationshipsWritten": result["relationshipsWritten"]
        }
    
    # =========================================================
    # CUSTOM SIMILARITY WITH CYPHER
    # =========================================================
    
    def compute_jaccard_cypher(
        self,
        label: str,
        relationship_type: str,
        node1_id: str,
        node2_id: str
    ) -> Optional[float]:
        """
        Compute Jaccard similarity using Cypher.
        
        This is useful when GDS is not available or for
        specific node pairs.
        """
        query = f"""
        MATCH (a:{label} {{id: $id1}})-[:{relationship_type}]->(common)
        WITH a, collect(common) as aNeighbors
        MATCH (b:{label} {{id: $id2}})-[:{relationship_type}]->(common)
        WITH a, aNeighbors, b, collect(common) as bNeighbors
        WITH a, b, aNeighbors, bNeighbors,
             [x IN aNeighbors WHERE x IN bNeighbors] as intersection,
             aNeighbors + [x IN bNeighbors WHERE NOT x IN aNeighbors] as union
        RETURN a.name as node1,
               b.name as node2,
               toFloat(size(intersection)) / size(union) as jaccard
        """
        
        with self.driver.session() as session:
            result = session.run(query, id1=node1_id, id2=node2_id)
            record = result.single()
            return record["jaccard"] if record else None
    
    def find_similar_patients_cypher(
        self,
        patient_id: str,
        min_shared: int = 1,
        limit: int = 10
    ) -> List[Dict[str, Any]]:
        """
        Find patients with similar diagnoses using Cypher.
        
        Computes similarity based on shared diseases.
        """
        query = """
        MATCH (p1:Patient {id: $id})-[:DIAGNOSED_WITH]->(d:Disease)
        WITH p1, collect(d) as diseases1
        MATCH (p2:Patient)-[:DIAGNOSED_WITH]->(d:Disease)
        WHERE p2 <> p1
        WITH p1, diseases1, p2, collect(d) as diseases2
        WITH p1, p2, diseases1, diseases2,
             [d IN diseases1 WHERE d IN diseases2] as shared
        WHERE size(shared) >= $min_shared
        WITH p1, p2,
             size(shared) as sharedCount,
             toFloat(size(shared)) / 
             size(diseases1 + [d IN diseases2 WHERE NOT d IN diseases1]) as jaccard
        RETURN p2.name as patient,
               p2.id as patientId,
               sharedCount,
               round(jaccard, 3) as similarity
        ORDER BY similarity DESC
        LIMIT $limit
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                id=patient_id,
                min_shared=min_shared,
                limit=limit
            )
            return [record.data() for record in result]
    
    # =========================================================
    # SIMILARITY-BASED RECOMMENDATIONS
    # =========================================================
    
    def recommend_medications(
        self,
        patient_id: str,
        top_k: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Recommend medications based on similar patients.
        
        1. Find similar patients (by diagnoses)
        2. Get medications from similar patients
        3. Filter out medications patient already takes
        4. Rank by frequency among similar patients
        """
        query = """
        // Find similar patients by shared diagnoses
        MATCH (p:Patient {id: $id})-[:DIAGNOSED_WITH]->(d:Disease)<-[:DIAGNOSED_WITH]-(similar:Patient)
        WHERE p <> similar
        WITH p, similar, count(d) as sharedDiseases
        ORDER BY sharedDiseases DESC
        LIMIT 10
        
        // Get medications from similar patients
        MATCH (similar)-[:TAKES]->(m:Medication)
        WHERE NOT (p)-[:TAKES]->(m)
        
        // Count and rank
        RETURN m.name as medication,
               m.drugClass as drugClass,
               count(DISTINCT similar) as similarPatientCount,
               avg(sharedDiseases) as avgSimilarity
        ORDER BY similarPatientCount DESC, avgSimilarity DESC
        LIMIT $limit
        """
        
        with self.driver.session() as session:
            result = session.run(query, id=patient_id, limit=top_k)
            return [record.data() for record in result]
    
    # =========================================================
    # CLEANUP
    # =========================================================
    
    def drop_graph(self, graph_name: str) -> bool:
        """Drop a projected graph."""
        if self.gds.graph.exists(graph_name)["exists"]:
            self.gds.graph.drop(graph_name)
            return True
        return False


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Similarity Algorithms - Demonstration")
    print("=" * 60)
    
    if not GDS_AVAILABLE:
        print("❌ GDS library not available")
        exit(1)
    
    try:
        analyzer = SimilarityAnalyzer()
        
        # Project graph
        print("\n📌 Project Bipartite Graph")
        print("-" * 40)
        
        try:
            result = analyzer.project_bipartite_graph(
                "patient_similarity",
                "Patient",
                "Disease",
                "DIAGNOSED_WITH"
            )
            print(f"   Nodes: {result['nodeCount']}")
            print(f"   Relationships: {result['relationshipCount']}")
            
            # Compute similarity
            print("\n📌 Node Similarity (Jaccard)")
            print("-" * 40)
            
            similar_pairs = analyzer.node_similarity_stream(
                "patient_similarity",
                top_k=5,
                similarity_cutoff=0.1
            )
            
            for pair in similar_pairs[:5]:
                print(f"   Node {pair['node1']} - Node {pair['node2']}: {pair['similarity']:.3f}")
            
            # Cleanup
            analyzer.drop_graph("patient_similarity")
            
        except Exception as e:
            print(f"   Graph projection failed: {e}")
        
        # Cypher-based similarity
        print("\n📌 Cypher-Based Similarity")
        print("-" * 40)
        
        similar_patients = analyzer.find_similar_patients_cypher(
            "P001",
            min_shared=1,
            limit=5
        )
        
        for p in similar_patients:
            print(f"   {p['patient']}: {p['similarity']:.3f} ({p['sharedCount']} shared)")
        
        print("\n" + "=" * 60)
        print("✅ Similarity algorithms demonstration complete!")
        print("=" * 60)
        
    except Exception as e:
        print(f"Error: {e}")
